from asyncore import write
from concurrent.futures import thread
import csv
import enum
from ipaddress import summarize_address_range
import math
from operator import truediv
from random import randint
import string
from time import sleep
from graphics import Maze
import threading
from gui import GUI


class Qlearning:
    def __init__(self, maze_, reward_, qtable_) -> None:

        self.maze = maze_
        self.reward = reward_
        self.qtable = qtable_

        self.gui = GUI(self.maze)

        self.alpha = 0.01
        self.gamma = 0.5

        self.border = -0.85

        self.steps = 0
        self.poll = 0
        # 迷宫的尺寸
        self.size = len(self.maze)

        # 设置起点、终点
        self.state = (0, 0)
        self.end = (self.size-1, self.size-1)

    def action(self) -> None:

        # 找到当前位置对应的qtalbe项
        x, y = self.state
        directions = self.qtable[x][y]
        q_max = max(directions)

        q_max_index = []
        for i in range(len(directions)):
            if directions[i] == q_max:
                q_max_index.append(i)

        direct = q_max_index[randint(0, len(q_max_index)-1)]

        # if direct == 0:
        #     print('↑', x, y)
        # elif direct == 1:
        #     print('↓', x, y)
        # elif direct == 2:
        #     print('←', x, y)
        # elif direct == 3:
        #     print('→', x, y)

        #  更新步数
        self.steps += 1

        # 判断边界
        if direct == 0 and x == 0 or direct == 1 and x == self.size-1 or \
                direct == 2 and y == 0 or direct == 3 and y == self.size-1:
            # print('===============border')
            self.qtable[x][y][direct] += self.alpha * \
                (self.reward[x][y] + self.gamma *
                 self.border - self.qtable[x][y][direct])
            return

        # 其他合法情况
        if direct == 0:
            nx, ny = x-1, y
        elif direct == 1:
            nx, ny = x+1, y
        elif direct == 2:
            nx, ny = x, y-1
        elif direct == 3:
            nx, ny = x, y+1

        # 更新qtable
        self.qtable[x][y][direct] += self.alpha * \
            (self.reward[nx][ny] + self.gamma *
             max(self.qtable[nx][ny]) - self.qtable[x][y][direct])
        # 如果下一步是墙，不更新位置
        if self.maze[nx][ny] != 1:
            self.state = nx, ny

            if direct == 0:
                self.gui.UP()
            elif direct == 1:
                self.gui.DOWN()
            elif direct == 2:
                self.gui.LEFT()
            elif direct == 3:
                self.gui.RIGHT()

        # else:
            # print('---------------wall')

    def isEnd(self) -> bool:
        return self.state == self.end

    def reset(self):
        self.state = (0, 0)
        self.steps = 0
        self.gui.reset()

    def writeQtable(self, filename) -> None:
        writeCSV(filename, self.qtable)

    def train(self, loops) -> None:
        for i in range(loops):
            while True:
                if self.isEnd():
                    print('次数:', self.poll, '步数:', self.steps)
                    self.reset()
                    self.poll += 1
                    break
                self.action()
        self.writeQtable('qtable.csv')


# 读csv文件
def readCSV(filename, type):
    arr = []
    tmp = []
    with open(filename, 'r') as f:
        for line in f:
            linelist = line.split(',')
            for index, item in enumerate(linelist):
                if type == 'float':
                    tmp.append(float(item))
                elif type == 'int':
                    tmp.append(int(item))
                else:
                    return []
            arr.append(tmp)
            tmp = []
        return arr


# 写csv文件
def writeCSV(filename, data):

    with open(filename, 'w') as f:
        f_csv = csv.writer(f)
        for i in range(len(data)):
            for j in range(len(data[0])):
                f_csv.writerow(data[i][j])


if __name__ == "__main__":

    # 初始化maze矩阵
    maze = readCSV('maze_map.csv', 'int')

    # 初始化reward矩阵
    reward = readCSV('reward.csv', 'float')

    # 迷宫尺寸
    size = len(maze)

    # 初始化qtable为0
    qtable = [[[0 for x in range(4)] for y in range(size)]
              for z in range(size)]

    qlearning = Qlearning(maze, reward, qtable)

    t1 = threading.Thread(target=qlearning.train, args=[500])

    qlearning.gui.renderMaze()

    t1.start()

    qlearning.gui.mainloop()
